from keras import layers
from keras.layers import Input, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D,concatenate
from keras.layers import Conv1D, MaxPooling1D, Dropout
from keras.models import Model




def TestModel(state_shape,n_actions):
    """
    Implementation of the HappyModel.

    Arguments:
    input_shape -- shape of the images of the dataset

    Returns:
    model -- a Model() instance in Keras
    """

    ### START CODE HERE ###
    # Feel free to use the suggested outline in the text above to get started, and run through the whole
    # exercise (including the later portions of this notebook) once. The come back also try out other
    # network architectures as well.
    # Define the input placeholder as a tensor with shape input_shape. Think of this as your input image!
    S_input = Input(state_shape ,name="state_input")

    # S1 = S_input[:, :, :-1]
    # S2 = S_input[:, :, -1]

    X = Conv1D(16,kernel_size=3,strides=1,padding='valid')(S_input)
    X = Conv1D(32,kernel_size=3,strides=1,padding='valid')(X)
    X = BatchNormalization()(X)
    X = Activation('relu')(X)
    X = Flatten()(X)
    X = Dense(100,'tanh')(X)
    Y = Dense(n_actions,activation='softmax')(X)
    model = Model(inputs=[S_input], outputs=Y, name="ScheduleModel")


    return model

testmodel = TestModel((12,12+20))

testmodel.summary()